import pymongo
import os
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# MongoDB 配置
MONGO_CONFIG = {
    'host': os.environ.get('MONGO_HOST', '116.63.51.45'),
    'port': int(os.environ.get('MONGO_PORT', 27017)),
    'username': os.environ.get('MONGO_USERNAME', 'root'),
    'password': os.environ.get('MONGO_PASSWORD', 'Zmzzmz010627!'),
    'auth_db': os.environ.get('MONGO_AUTH_DB', 'admin'),
    'db': os.environ.get('MONGO_DB', 'kube-replayer')
}

# MongoDB 客户端
_client = None

def get_mongo_client():
    """获取 MongoDB 客户端连接"""
    global _client
    if _client is None:
        try:
            # 创建 MongoDB 客户端
            _client = pymongo.MongoClient(
                host=MONGO_CONFIG['host'],
                port=MONGO_CONFIG['port'],
                username=MONGO_CONFIG['username'],
                password=MONGO_CONFIG['password'],
                authSource=MONGO_CONFIG['auth_db']
            )
            # 测试连接
            _client.admin.command('ping')
            logger.info("MongoDB 连接成功")
        except Exception as e:
            logger.error(f"MongoDB 连接失败: {e}")
            raise
    return _client

def get_db():
    """获取 MongoDB 数据库"""
    client = get_mongo_client()
    return client[MONGO_CONFIG['db']]

def init_db():
    """初始化数据库，创建集合和索引"""
    try:
        db = get_db()
        
        # 创建集合（如果不存在）
        if 'record_services' not in db.list_collection_names():
            db.create_collection('record_services')
        if 'request_records' not in db.list_collection_names():
            db.create_collection('request_records')
        if 'request_definitions' not in db.list_collection_names():
            db.create_collection('request_definitions')
        if 'intercept_services' not in db.list_collection_names():
            db.create_collection('intercept_services')
        if 'intercept_request_records' not in db.list_collection_names():
            db.create_collection('intercept_request_records')
        
        # 创建索引
        db.record_services.create_index([('service_name', pymongo.ASCENDING)], unique=True)
        db.request_records.create_index([('service_id', pymongo.ASCENDING)])
        db.request_records.create_index([('created_at', pymongo.DESCENDING)])
        
        db.intercept_services.create_index([('service_name', pymongo.ASCENDING)], unique=True)
        db.intercept_request_records.create_index([('intercept_service_id', pymongo.ASCENDING)])
        db.intercept_request_records.create_index([
            ('intercept_service_id', pymongo.ASCENDING),
            ('request_method', pymongo.ASCENDING),
            ('request_url', pymongo.ASCENDING)
        ], unique=True)
        
        logger.info("MongoDB 初始化成功")
    except Exception as e:
        logger.error(f"MongoDB 初始化失败: {e}")
        raise

def close_connection():
    """关闭 MongoDB 连接"""
    global _client
    if _client is not None:
        _client.close()
        _client = None
        logger.info("MongoDB 连接已关闭")

# 辅助函数：将 MongoDB 文档转换为 snake_case
def to_snake_case(name):
    """将驼峰命名转换为蛇形命名"""
    import re
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()

def convert_to_snake_case(doc):
    """递归将文档的键转换为蛇形命名"""
    if isinstance(doc, dict):
        new_doc = {}
        for k, v in doc.items():
            if k == '_id':
                new_doc[k] = v
            else:
                new_doc[to_snake_case(k)] = convert_to_snake_case(v)
        return new_doc
    elif isinstance(doc, list):
        return [convert_to_snake_case(item) for item in doc]
    else:
        return doc

if __name__ == "__main__":
    # 当直接运行此脚本时，初始化数据库
    init_db()
    close_connection()
